package study.eggs.async.mdc;

import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.Callable;

@Slf4j
public final class MdcUtils {

    /**
     * 包装 Runnable
     *
     * @param runnable
     * @return Runnable
     */
    public static Runnable wrap(final Runnable runnable) {
        Map<String, String> context = MDC.getCopyOfContextMap();
        return () -> {
            //TODO 思考：此处获取 MDC 上下文为空
            //Map<String, String> context = MDC.getCopyOfContextMap();
            if (context == null) {
                MDC.clear();
            } else {
                //将主线程的Map设置到当前线程中
                MDC.setContextMap(context);
            }
            try {
                runnable.run();
            } finally {
                MDC.clear();
            }
        };
    }

    /**
     * 包装 Callable
     *
     * @param callable
     * @return Callable
     */
    public static <T> Callable<T> wrap(final Callable<T> callable) {
        return () -> {
            Map<String, String> context = MDC.getCopyOfContextMap();
            if (context == null) {
                MDC.clear();
            } else {
                //将主线程的Map设置到当前线程中
                MDC.setContextMap(context);
            }
            try {
                return callable.call();
            } finally {
                MDC.clear();
            }
        };
    }
}
